// Yue Marvin Tao
// 3/27/2022

package mult 

import chisel3._
import chisel3.util._

class BMInputBundle(val bits: Int) extends Bundle {
  val load = Input(Bool())
  val A = Input(UInt(bits.W))
  val B = Input(UInt(bits.W))
}

class BMOutputBundle(val bits: Int) extends Bundle {
  val P = Output(UInt((bits * 2).W))
  val Q = Output(UInt(1.W))
  val c = Output(UInt(bits.W))
  val comp = Output(UInt(bits.W))
  val shift = Output(Bool())
}

class BoothMul(val bits: Int) extends Module {
  val in = IO(new BMInputBundle(bits))
  val out = IO(new BMOutputBundle(bits))

  val Acomp = Wire(UInt(bits.W)) 

  val prod = RegNext(VecInit(Seq.fill(2 * bits)(0.U(1.W))))
  val P    = prod(0)
  val Q    = RegNext(0.U(1.W))
  val carry = Wire(Vec(bits + 1, UInt(1.W))) 
  val shift = RegNext(false.B)
  val count = RegNext(0.U(bits.W))

  // Initializing 'A and carry flags
  Acomp := (1 << bits).U - in.A
  for (i <- 0 until bits + 1) {
    carry(i) := 0.U
  }

  when(in.load) {
    // load B (multiplier) on low bits (except the last) of product
    Q := 0.U
    for (i <- 0 until bits) {
      prod(i) := in.B(i)
      prod(i + bits) := 0.U
    }
  }.elsewhen(count < bits.U){
    // add A or 'A up to high bits of product 
    when(shift) {

      Q := prod(0)
      for (i <- 0 until (2 * bits - 1)) {
        prod(i) := prod(i+1)
      }
      prod(2*bits-1) := prod(2*bits-1)
      count := count + 1.U
      shift := false.B

    }.elsewhen(Q === 0.U && P === 1.U) {

      for (i <- 0 until bits) {
        val p = prod(i + bits)
        val m = Acomp(i)
        val c = carry(i)
        val cn = carry(i+1)
        p := p ^ m ^ c 
        cn := (p & m) | (m & c) | (c & p) 
      }
      for (i <- 0 until bits) {
        prod(i) := prod(i)
      }
      shift := true.B

    }.elsewhen(Q === 1.U && P === 0.U) {

      for (i <- 0 until bits) {
        val p = prod(i + bits)
        val m = in.A(i)
        val c = carry(i)
        val cn = carry(i+1)
        p := p ^ m ^ c 
        cn := (p & m) | (m & c) | (c & p) 
      }
      for (i <- 0 until bits) {
        prod(i) := prod(i)
      }
      shift := true.B

    }.otherwise {
      
      prod := prod
      Q := Q
      shift := true.B
    
    }

  }

  out.P := prod.asUInt
  out.Q := Q
  out.c := count
  out.comp := Acomp
  out.shift := shift
}